import numpy as np

class MarketStateAnalyzer:
    def __init__(self):
        pass
    
    def calculate_adx(self, data, period=14):
        """计算ADX指标"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            high = data['high'].values
            low = data['low'].values
            close = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            high = [d['high'] for d in data]
            low = [d['low'] for d in data]
            close = [d['close'] for d in data]
            length = len(data)
        
        if length < period + 1:
            return [None] * length
        
        # 计算+DM和-DM
        plus_dm = [0.0] * length
        minus_dm = [0.0] * length
        
        for i in range(1, length):
            up_move = high[i] - high[i-1]
            down_move = low[i-1] - low[i]
            
            if up_move > down_move and up_move > 0:
                plus_dm[i] = up_move
            if down_move > up_move and down_move > 0:
                minus_dm[i] = down_move
        
        # 计算+DI和-DI
        tr = [0.0] * length
        for i in range(1, length):
            tr[i] = max(
                high[i] - low[i],
                abs(high[i] - close[i-1]),
                abs(low[i] - close[i-1])
            )
        
        plus_di = [None] * length
        minus_di = [None] * length
        dx = [None] * length
        adx = [None] * length
        
        # 计算初始+DI和-DI
        plus_dm_sum = sum(plus_dm[1:period+1])
        minus_dm_sum = sum(minus_dm[1:period+1])
        tr_sum = sum(tr[1:period+1])
        
        if tr_sum != 0:
            plus_di[period] = 100 * (plus_dm_sum / tr_sum)
            minus_di[period] = 100 * (minus_dm_sum / tr_sum)
        
        # 计算后续值
        for i in range(period + 1, length):
            plus_dm_sum = plus_dm_sum - (plus_dm_sum / period) + plus_dm[i]
            minus_dm_sum = minus_dm_sum - (minus_dm_sum / period) + minus_dm[i]
            tr_sum = tr_sum - (tr_sum / period) + tr[i]
            
            if tr_sum != 0:
                plus_di[i] = 100 * (plus_dm_sum / tr_sum)
                minus_di[i] = 100 * (minus_dm_sum / tr_sum)
                
                di_diff = abs(plus_di[i] - minus_di[i])
                di_sum = plus_di[i] + minus_di[i]
                if di_sum != 0:
                    dx[i] = 100 * (di_diff / di_sum)
        
        # 计算ADX
        if period < length:
            adx[period] = dx[period]
            
        for i in range(period + 1, length):
            if dx[i] is not None:
                if adx[i-1] is not None:
                    adx[i] = (adx[i-1] * (period - 1) + dx[i]) / period
                else:
                    # 计算前period个dx的平均值
                    dx_values = [dx[j] for j in range(period, i) if dx[j] is not None]
                    if len(dx_values) > 0:
                        adx[i] = sum(dx_values) / len(dx_values)
                    else:
                        adx[i] = None
        
        return adx
    
    def calculate_atr(self, data, period=14):
        """计算ATR指标"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            high = data['high'].values
            low = data['low'].values
            close = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            high = [d['high'] for d in data]
            low = [d['low'] for d in data]
            close = [d['close'] for d in data]
            length = len(data)
        
        if length < period + 1:
            return [None] * length
        
        tr = [0.0] * length
        for i in range(1, length):
            tr[i] = max(
                high[i] - low[i],
                abs(high[i] - close[i-1]),
                abs(low[i] - close[i-1])
            )
        
        atr = [None] * length
        # 计算初始ATR
        sum_tr = sum(tr[1:period+1])
        atr[period] = sum_tr / period
        
        # 计算后续ATR
        for i in range(period + 1, length):
            atr[i] = (atr[i-1] * (period - 1) + tr[i]) / period
        
        return atr
    
    def calculate_bollinger_bands(self, data, period=20, multiplier=2):
        """计算布林带"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            closes = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            closes = [d['close'] for d in data]
            length = len(data)
        
        if length < period:
            return {
                'upper': [None] * length,
                'middle': [None] * length,
                'lower': [None] * length
            }
        
        middle = [None] * length
        upper = [None] * length
        lower = [None] * length
        
        for i in range(period - 1, length):
            mean = sum(closes[i-period+1:i+1]) / period
            variance = sum((close - mean) ** 2 for close in closes[i-period+1:i+1]) / period
            std_dev = variance ** 0.5
            
            middle[i] = mean
            upper[i] = mean + (std_dev * multiplier)
            lower[i] = mean - (std_dev * multiplier)
        
        return {'upper': upper, 'middle': middle, 'lower': lower}
    
    def analyze_high_low_structure(self, data, period=20):
        """分析高低点结构 - 重写版本"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            high = data['high'].values
            low = data['low'].values
            close = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            high = [d['high'] for d in data]
            low = [d['low'] for d in data]
            close = [d['close'] for d in data]
            length = len(data)
        
        if length < period:
            return {
                'higher_highs': False,
                'lower_lows': False,
                'higher_lows': False,
                'lower_highs': False,
                'trend_structure': 'neutral',
                'structure_strength': 0,
                'swing_points': []
            }
        
        # 找出局部高点和低点（摆动点）
        highs = []  # (index, price, type)
        lows = []   # (index, price, type)
        
        # 使用更严格的高低点识别算法
        for i in range(2, length - 2):
            # 识别高点：必须比左右两边都高
            if (high[i] > high[i-1] and high[i] > high[i-2] and 
                high[i] > high[i+1] and high[i] > high[i+2]):
                highs.append((i, high[i], 'resistance'))
            
            # 识别低点：必须比左右两边都低
            if (low[i] < low[i-1] and low[i] < low[i-2] and 
                low[i] < low[i+1] and low[i] < low[i+2]):
                lows.append((i, low[i], 'support'))
        
        # 分析高低点序列趋势
        higher_highs = False
        lower_lows = False
        higher_lows = False
        lower_highs = False
        
        # 分析最近的高点趋势（至少需要3个点才能判断趋势）
        if len(highs) >= 3:
            recent_highs = highs[-3:]  # 取最近3个高点
            # 检查是否形成更高的高点
            higher_highs = (recent_highs[0][1] < recent_highs[1][1] < recent_highs[2][1])
            # 检查是否形成更低的高点
            lower_highs = (recent_highs[0][1] > recent_highs[1][1] > recent_highs[2][1])
        
        # 分析最近的低点趋势（至少需要3个点才能判断趋势）
        if len(lows) >= 3:
            recent_lows = lows[-3:]  # 取最近3个低点
            # 检查是否形成更高的低点
            higher_lows = (recent_lows[0][1] < recent_lows[1][1] < recent_lows[2][1])
            # 检查是否形成更低的低点
            lower_lows = (recent_lows[0][1] > recent_lows[1][1] > recent_lows[2][1])
        
        # 判断趋势结构 - 修复逻辑
        trend_structure = 'neutral'
        structure_strength = 0
        
        # 上涨趋势：更高的高点 + 更高的低点
        if higher_highs and higher_lows:
            trend_structure = 'uptrend'
            structure_strength = 0.8
        # 下跌趋势：更低的高点 + 更低的低点
        elif lower_highs and lower_lows:
            trend_structure = 'downtrend'
            structure_strength = 0.8
        # 潜在上涨趋势：更高的高点但低点不明显
        elif higher_highs and not lower_lows:
            trend_structure = 'potential_uptrend'
            structure_strength = 0.5
        # 潜在下跌趋势：更低的低点但高点不明显
        elif lower_lows and not higher_highs:
            trend_structure = 'potential_downtrend'
            structure_strength = 0.5
        # 震荡结构：高低点混乱
        elif len(highs) >= 2 and len(lows) >= 2:
            # 检查是否在区间内震荡
            high_range = max(h[1] for h in highs) - min(h[1] for h in highs)
            low_range = max(l[1] for l in lows) - min(l[1] for l in lows)
            price_range = max(close) - min(close)
            
            if high_range < price_range * 0.3 and low_range < price_range * 0.3:
                trend_structure = 'ranging'
                structure_strength = 0.6
        
        # 计算结构强度（基于摆动点的数量和清晰度）
        total_swing_points = len(highs) + len(lows)
        if total_swing_points >= 6:  # 有足够的摆动点
            structure_strength = min(1.0, structure_strength + 0.2)
        elif total_swing_points < 4:  # 摆动点不足
            structure_strength = max(0.2, structure_strength - 0.3)
        
        # 收集重要的摆动点信息
        swing_points = []
        # 添加最近的重要高点和低点
        for i, (idx, price, point_type) in enumerate(highs[-4:]):
            swing_points.append({
                'index': idx,
                'price': price,
                'type': point_type,
                'importance': len(highs) - i  # 越近的点越重要
            })
        
        for i, (idx, price, point_type) in enumerate(lows[-4:]):
            swing_points.append({
                'index': idx,
                'price': price,
                'type': point_type,
                'importance': len(lows) - i  # 越近的点越重要
            })
        
        # 按重要性排序
        swing_points.sort(key=lambda x: x['importance'], reverse=True)
        
        return {
            'higher_highs': higher_highs,
            'lower_lows': lower_lows,
            'higher_lows': higher_lows,
            'lower_highs': lower_highs,
            'trend_structure': trend_structure,
            'structure_strength': structure_strength,
            'swing_points': swing_points[:6]  # 返回最重要的6个摆动点
        }
    
    def analyze_ma_alignment(self, data, short_period=5, long_period=20):
        """分析均线排列"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            closes = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            closes = [d['close'] for d in data]
            length = len(data)
        
        if length < long_period:
            return {
                'ma_short': None,
                'ma_long': None,
                'alignment': 'neutral'
            }
        
        # 计算短期和长期均线
        ma_short = sum(closes[-short_period:]) / short_period
        ma_long = sum(closes[-long_period:]) / long_period
        
        # 判断均线排列
        alignment = 'neutral'
        if ma_short > ma_long:
            alignment = 'bullish'  # 多头排列
        elif ma_short < ma_long:
            alignment = 'bearish'  # 空头排列
        
        return {
            'ma_short': ma_short,
            'ma_long': ma_long,
            'alignment': alignment
        }
    
    def calculate_volatility(self, data, period=14):
        """计算波动率"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            closes = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            closes = [d['close'] for d in data]
            length = len(data)
        
        if length < period:
            return [None] * length
        
        volatility = [None] * length
        
        for i in range(period, length):
            # 计算周期内的价格变化百分比
            changes = []
            for j in range(i - period + 1, i + 1):
                if j > 0:
                    change = abs((closes[j] - closes[j-1]) / closes[j-1] * 100)
                    changes.append(change)
            
            if changes:
                volatility[i] = sum(changes) / len(changes)
        
        return volatility
    
    def determine_market_regime(self, data, adx_values, atr_values, bb_bands, high_low_structure, ma_alignment, volatility_values):
        """综合判断市场状态（趋势vs震荡）"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            close = data['close'].values
            length = len(data)
        else:
            # 字典列表格式
            close = [d['close'] for d in data]
            length = len(data)
        
        if length < 20:
            return {
                'regime': 'neutral',
                'confidence': 0.5,
                'reasons': ['数据不足']
            }
        
        last_index = length - 1
        
        # 获取各指标的最新值
        adx = adx_values[last_index] if adx_values[last_index] is not None else None
        atr = atr_values[last_index] if atr_values[last_index] is not None else None
        volatility = volatility_values[last_index] if volatility_values[last_index] is not None else None
        
        # 获取布林带值
        bb_upper = bb_bands['upper'][last_index] if bb_bands['upper'][last_index] is not None else None
        bb_middle = bb_bands['middle'][last_index] if bb_bands['middle'][last_index] is not None else None
        bb_lower = bb_bands['lower'][last_index] if bb_bands['lower'][last_index] is not None else None
        
        # 获取价格信息
        current_price = close[last_index]
        
        # 初始化评分
        trend_score = 0
        oscillation_score = 0
        reasons = []
        
        # 1. ADX判断
        if adx is not None:
            if adx > 25:
                trend_score += 2
                reasons.append(f"ADX({adx:.1f})>25表明有趋势")
            elif adx < 20:
                oscillation_score += 2
                reasons.append(f"ADX({adx:.1f})<20表明无趋势")
        
        # 2. ATR判断
        if atr is not None and current_price is not None:
            atr_ratio = (atr / current_price) * 100  # ATR占价格的百分比
            if atr_ratio > 1.5:
                trend_score += 1
                reasons.append(f"ATR占比({atr_ratio:.2f}%)>1.5%表明波动大，可能是趋势行情")
            elif atr_ratio < 1:
                oscillation_score += 1
                reasons.append(f"ATR占比({atr_ratio:.2f}%)<1%表明波动小，可能是震荡行情")
        
        # 3. 布林带判断
        if all(v is not None for v in [bb_upper, bb_middle, bb_lower, current_price]):
            # 判断价格是否贴轨
            upper_distance = abs(current_price - bb_upper)
            lower_distance = abs(current_price - bb_lower)
            middle_distance = abs(current_price - bb_middle)
            
            # 判断布林带开口
            bb_width = (bb_upper - bb_lower) / bb_middle * 100  # 布林带宽度百分比
            
            if upper_distance < (bb_upper - bb_middle) * 0.2:  # 价格接近上轨
                trend_score += 1
                reasons.append("价格接近布林带上轨")
            elif lower_distance < (bb_middle - bb_lower) * 0.2:  # 价格接近下轨
                trend_score += 1
                reasons.append("价格接近布林带下轨")
            elif middle_distance < (bb_upper - bb_lower) * 0.2:  # 价格在中轨附近
                oscillation_score += 1
                reasons.append("价格在布林带中轨附近")
            
            if bb_width > 2:  # 布林带开口较大
                trend_score += 1
                reasons.append(f"布林带开口({bb_width:.2f}%)较大表明趋势行情")
            elif bb_width < 1:  # 布林带开口较小
                oscillation_score += 1
                reasons.append(f"布林带开口({bb_width:.2f}%)较小表明震荡行情")
        
        # 4. 高低点结构判断 - 更新逻辑
        structure_strength = high_low_structure.get('structure_strength', 0)
        
        if high_low_structure['trend_structure'] == 'uptrend':
            trend_score += int(2 * structure_strength)
            reasons.append(f"高点和低点持续抬高表明上涨趋势 (强度: {structure_strength:.1f})")
        elif high_low_structure['trend_structure'] == 'downtrend':
            trend_score += int(2 * structure_strength)
            reasons.append(f"高点和低点持续降低表明下跌趋势 (强度: {structure_strength:.1f})")
        elif high_low_structure['trend_structure'] == 'potential_uptrend':
            trend_score += int(1 * structure_strength)
            reasons.append(f"可能形成上涨趋势，高点抬高 (强度: {structure_strength:.1f})")
        elif high_low_structure['trend_structure'] == 'potential_downtrend':
            trend_score += int(1 * structure_strength)
            reasons.append(f"可能形成下跌趋势，低点降低 (强度: {structure_strength:.1f})")
        elif high_low_structure['trend_structure'] == 'ranging':
            oscillation_score += int(2 * structure_strength)
            reasons.append(f"价格在区间内震荡 (强度: {structure_strength:.1f})")
        else:  # neutral
            oscillation_score += 1
            reasons.append("高低点结构不明显，可能震荡")
        
        # 5. 均线排列判断
        if ma_alignment['alignment'] == 'bullish':
            trend_score += 1
            reasons.append("均线多头排列表明上涨趋势")
        elif ma_alignment['alignment'] == 'bearish':
            trend_score += 1
            reasons.append("均线空头排列表明下跌趋势")
        elif ma_alignment['alignment'] == 'neutral':
            oscillation_score += 1
            reasons.append("均线无明显排列表明震荡行情")
        
        # 6. 波动率判断
        if volatility is not None:
            if volatility > 1.5:
                trend_score += 1
                reasons.append(f"平均波动率({volatility:.2f}%)>1.5%表明趋势行情")
            elif volatility < 1:
                oscillation_score += 1
                reasons.append(f"平均波动率({volatility:.2f}%)<1%表明震荡行情")
        
        # 综合判断
        regime = 'neutral'
        confidence = 0.5
        
        if trend_score > oscillation_score:
            regime = 'trend'
            confidence = min(0.9, 0.5 + (trend_score - oscillation_score) * 0.1)
        elif oscillation_score > trend_score:
            regime = 'oscillating'
            confidence = min(0.9, 0.5 + (oscillation_score - trend_score) * 0.1)
        
        return {
            'regime': regime,
            'confidence': confidence,
            'reasons': reasons
        }
    
    def get_market_regime_description(self, regime_result):
        """获取市场状态描述"""
        regime = regime_result['regime']
        confidence = regime_result['confidence']
        
        if regime == 'trend':
            return f"趋势行情 (置信度: {confidence*100:.1f}%)"
        elif regime == 'oscillating':
            return f"震荡行情 (置信度: {confidence*100:.1f}%)"
        else:
            return f"中性状态 (置信度: {confidence*100:.1f}%)"
    
    def analyze_market_state(self, data):
        """完整分析市场状态"""
        # 确保数据是DataFrame格式
        if hasattr(data, 'iloc'):
            # DataFrame格式
            length = len(data)
        else:
            # 字典列表格式
            length = len(data)
        
        if not data or length < 20:
            return {
                'regime': 'neutral',
                'confidence': 0.5,
                'description': '数据不足，无法判断市场状态',
                'details': {}
            }
        
        # 计算各项指标
        adx_values = self.calculate_adx(data)
        atr_values = self.calculate_atr(data)
        bb_bands = self.calculate_bollinger_bands(data)
        high_low_structure = self.analyze_high_low_structure(data)
        ma_alignment = self.analyze_ma_alignment(data)
        volatility_values = self.calculate_volatility(data)
        
        # 综合判断市场状态
        regime_result = self.determine_market_regime(
            data, adx_values, atr_values, bb_bands, 
            high_low_structure, ma_alignment, volatility_values
        )
        
        # 获取描述
        description = self.get_market_regime_description(regime_result)
        
        # 返回详细信息
        details = {
            'adx': adx_values[-1] if adx_values[-1] is not None else 'N/A',
            'atr': atr_values[-1] if atr_values[-1] is not None else 'N/A',
            'volatility': volatility_values[-1] if volatility_values[-1] is not None else 'N/A',
            'high_low_structure': high_low_structure['trend_structure'],
            'ma_alignment': ma_alignment['alignment'],
            'bb_width': 'N/A',
            'price_position': 'N/A'
        }
        
        # 计算布林带相关信息
        last_index = length - 1
        if (bb_bands['upper'][last_index] is not None and 
            bb_bands['middle'][last_index] is not None and 
            bb_bands['lower'][last_index] is not None):
            bb_width = (bb_bands['upper'][last_index] - bb_bands['lower'][last_index]) / bb_bands['middle'][last_index] * 100
            details['bb_width'] = f"{bb_width:.2f}%"
            
            # 计算价格位置
            current_price = close[last_index] if 'close' in locals() else data[last_index]['close']
            upper = bb_bands['upper'][last_index]
            middle = bb_bands['middle'][last_index]
            lower = bb_bands['lower'][last_index]
            
            if current_price > middle:
                position = ((current_price - middle) / (upper - middle)) * 100
                details['price_position'] = f"布林带上部 {position:.1f}%"
            else:
                position = ((middle - current_price) / (middle - lower)) * 100
                details['price_position'] = f"布林带下部 {position:.1f}%"
        
        return {
            'regime': regime_result['regime'],
            'confidence': regime_result['confidence'],
            'description': description,
            'details': details,
            'reasons': regime_result['reasons']
        }